#!/usr/bin/env python3
"""
RAG文档摄取系统快速测试脚本

这个脚本帮助您快速测试RAG文档摄取系统，无需手动准备大量文档。
"""

import os
import sys
from rag_document_ingestion import ingest_documents

def create_sample_documents():
    """创建示例文档用于测试"""
    
    # 确保data目录存在
    os.makedirs("data", exist_ok=True)
    
    # 创建示例Markdown文档（可以作为文本文件处理）
    sample_content = """
# 人工智能与机器学习概述

## 什么是人工智能？

人工智能（Artificial Intelligence, AI）是计算机科学的一个分支，致力于创建能够执行通常需要人类智能的任务的系统。

## 机器学习的基本概念

机器学习是人工智能的一个子集，它使用算法和统计模型，使计算机系统能够从数据中自动学习和改进，而无需明确编程。

### 主要类型：

1. **监督学习**：使用标记数据训练模型
2. **无监督学习**：从未标记数据中发现模式
3. **强化学习**：通过与环境交互学习最优行为

## 深度学习

深度学习是机器学习的一个分支，使用多层神经网络来学习数据的复杂模式。

### 应用领域：
- 计算机视觉
- 自然语言处理
- 语音识别
- 推荐系统

## 大语言模型

大语言模型（Large Language Models, LLMs）是近年来AI领域的重要突破，能够理解和生成人类语言。

### 代表性模型：
- GPT系列（OpenAI）
- BERT（Google）
- 文心一言（百度）
- 通义千问（阿里巴巴）

## RAG技术

检索增强生成（Retrieval-Augmented Generation）结合了信息检索和文本生成，是当前AI应用的重要技术方向。
"""
    
    # 保存为文本文件
    with open("data/ai_overview.txt", "w", encoding="utf-8") as f:
        f.write(sample_content)
    
    # 创建另一个示例文档
    tech_content = """
技术栈学习指南

编程语言选择
Python: 适合数据科学、AI、Web开发
JavaScript: 前端开发必备，后端也可用
Java: 企业级应用开发
Go: 高性能后端服务
Rust: 系统级编程，内存安全

数据库技术
关系型数据库：MySQL, PostgreSQL
非关系型数据库：MongoDB, Redis
向量数据库：Chroma, Pinecone, Qdrant

机器学习框架
TensorFlow: Google开发，功能全面
PyTorch: Facebook开发，研究友好
Scikit-learn: 传统机器学习算法
Hugging Face: 预训练模型库

开发工具
版本控制：Git, GitHub
代码编辑器：VSCode, PyCharm
虚拟环境：conda, venv
包管理：pip, npm, yarn

云服务平台
AWS: 亚马逊云服务
Azure: 微软云服务
阿里云: 国内领先云服务
腾讯云: 游戏和社交领域优势

学习建议
1. 选择一门主要编程语言深入学习
2. 掌握基础的数据结构和算法
3. 实践项目驱动学习
4. 关注开源社区和技术博客
5. 参与技术交流和分享
"""
    
    with open("data/tech_guide.txt", "w", encoding="utf-8") as f:
        f.write(tech_content)
    
    print("✅ 示例文档创建完成！")
    print("📂 创建的文件：")
    print("   - data/ai_overview.txt")
    print("   - data/tech_guide.txt")
    print("   - data/notes.txt (已存在)")

def main():
    """主函数"""
    print("🚀 RAG文档摄取系统测试")
    print("=" * 50)
    
    # BGE-M3模型信息
    print("🔧 嵌入模型信息：")
    print("   - 模型: BAAI/bge-m3")
    print("   - 运行方式: 本地运行，无需API密钥")
    print("   - 首次运行会自动下载模型文件（约2.3GB）")
    print()
    
    # 创建示例文档
    print("📝 创建示例文档...")
    create_sample_documents()
    print()
    
    # 检查现有文档
    test_files = []
    data_dir = "./data"
    
    if os.path.exists(data_dir):
        for file in os.listdir(data_dir):
            if file.endswith(('.txt', '.pdf', '.docx')):
                test_files.append(os.path.join(data_dir, file))
    
    if not test_files:
        print("❌ 在data目录中未找到支持的文档文件")
        return False
    
    print(f"📚 找到 {len(test_files)} 个文档文件：")
    for file in test_files:
        print(f"   - {file}")
    print()
    
    try:
        # 执行文档摄取
        print("🔄 开始文档摄取流程...")
        vector_store = ingest_documents(
            file_paths=test_files,
            persist_directory="./chroma_db"
        )
        
        print()
        print("🎉 文档摄取完成！")
        print(f"📊 向量数据库信息：")
        print(f"   - 位置: ./chroma_db")
        print(f"   - 集合: rag_documents")
        
        # 测试检索功能
        print()
        print("🔍 测试检索功能...")
        test_queries = [
            "什么是RAG系统？",
            "机器学习的类型有哪些？",
            "推荐的编程语言",
            "向量数据库"
        ]
        
        for query in test_queries:
            print(f"\n查询: {query}")
            results = vector_store.similarity_search(query, k=2)
            
            if results:
                print(f"找到 {len(results)} 个相关结果：")
                for i, result in enumerate(results):
                    preview = result.page_content[:100].replace('\n', ' ')
                    source = result.metadata.get('source_file', 'unknown')
                    print(f"  {i+1}. [{os.path.basename(source)}] {preview}...")
            else:
                print("  未找到相关结果")
        
        print()
        print("✅ 测试完成！系统运行正常。")
        print()
        print("📖 后续步骤：")
        print("1. 将您的文档放入data目录")
        print("2. 运行 python rag_document_ingestion.py 处理您的文档")
        print("3. 使用vector_store.similarity_search()进行检索")
        
        return True
        
    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")
        print()
        print("🔧 故障排除建议：")
        print("1. 检查网络连接（首次下载模型需要）")
        print("2. 确认已安装sentence-transformers包")
        print("3. 确保有足够的内存和存储空间")
        print("4. 查看详细错误信息进行调试")
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
